{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 06, case 7b: Stokes problem with Neumann control\n",
    "\n",
    "In this tutorial we solve the optimal control problem\n",
    "\n",
    "$$\\min J(y, u) = \\frac{1}{2} \\int_{\\Gamma_{obs}} (v - v_d)^2 dx + \\frac{\\alpha_1}{2} \\int_{\\Gamma_C} |\\nabla_{\\mathbf{t}} u|^2 ds + \\frac{\\alpha_2}{2} \\int_{\\Gamma_C} |u|^2 ds$$\n",
    "s.t.\n",
    "$$\\begin{cases}\n",
    "- \\nu \\Delta v + \\nabla p = f       & \\text{in } \\Omega\\\\\n",
    "             \\text{div} v = 0       & \\text{in } \\Omega\\\\\n",
    "                        v = g       & \\text{on } \\Gamma_{in}\\\\\n",
    "                        v = 0       & \\text{on } \\Gamma_{w}\\\\\n",
    "   p n - \\nu \\partial_n v = u       & \\text{on } \\Gamma_{C}\n",
    "\\end{cases}$$\n",
    "\n",
    "where\n",
    "$$\\begin{align*}\n",
    "& \\Omega                      & \\text{unit square}\\\\\n",
    "& \\Gamma_{in}                 & \\text{has boundary id 1}\\\\\n",
    "& \\Gamma_{w}                  & \\text{has boundary id 2}\\\\\n",
    "& \\Gamma_{C}                  & \\text{has boundary id 3}\\\\\n",
    "& \\Gamma_{obs}                & \\text{has boundary id 4}\\\\\n",
    "& u \\in [L^2(\\Gamma_C)]^2     & \\text{control variable}\\\\\n",
    "& v \\in [H^1(\\Omega)]^2       & \\text{state velocity variable}\\\\\n",
    "& p \\in L^2(\\Omega)           & \\text{state pressure variable}\\\\\n",
    "& \\alpha_1, \\alpha_2 > 0      & \\text{penalization parameters}\\\\\n",
    "& v_d                         & \\text{desired state}\\\\\n",
    "& f                           & \\text{forcing term}\\\\\n",
    "& g                           & \\text{inlet profile}\\\\\n",
    "\\end{align*}$$\n",
    "using an adjoint formulation solved by a one shot approach"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from mpi4py import MPI\n",
    "from petsc4py import PETSc\n",
    "from ufl import (as_vector, div, FacetNormal, grad, inner, Measure, replace, SpatialCoordinate,\n",
    "                 TestFunction, TrialFunction)\n",
    "from dolfinx import Constant, DirichletBC, Function, FunctionSpace, VectorFunctionSpace\n",
    "from dolfinx.cpp.mesh import GhostMode\n",
    "from dolfinx.fem import locate_dofs_topological\n",
    "from dolfinx.io import XDMFFile\n",
    "from dolfinx.plot import create_vtk_topology\n",
    "from multiphenicsx.fem import (assemble_matrix_block, assemble_scalar, assemble_vector_block,\n",
    "                               BlockVecSubVectorWrapper, create_vector_block, DofMapRestriction)\n",
    "import pyvista"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if MPI.COMM_WORLD.size > 1:\n",
    "    mesh_ghost_mode = GhostMode.shared_facet  # shared_facet ghost mode is required by dS\n",
    "else:\n",
    "    mesh_ghost_mode = GhostMode.none\n",
    "with XDMFFile(MPI.COMM_WORLD, \"data/bifurcation.xdmf\", \"r\") as infile:\n",
    "    mesh = infile.read_mesh(mesh_ghost_mode)\n",
    "    mesh.topology.create_connectivity_all()\n",
    "    subdomains = infile.read_meshtags(mesh, name=\"subdomains\")\n",
    "    boundaries = infile.read_meshtags(mesh, name=\"boundaries\")\n",
    "boundaries_1 = boundaries.indices[boundaries.values == 1]\n",
    "boundaries_2 = boundaries.indices[boundaries.values == 2]\n",
    "boundaries_3 = boundaries.indices[boundaries.values == 3]\n",
    "boundaries_12 = boundaries.indices[np.isin(boundaries.values, (1, 2))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define associated measures\n",
    "dx = Measure(\"dx\")(subdomain_data=subdomains)\n",
    "ds = Measure(\"ds\")(subdomain_data=boundaries)\n",
    "dS = Measure(\"dS\")(subdomain_data=boundaries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Normal and tangent\n",
    "n = FacetNormal(mesh)\n",
    "t = as_vector([n[1], -n[0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dolfinx_to_pyvista_mesh(mesh):\n",
    "    num_cells = mesh.topology.index_map(mesh.topology.dim).size_local\n",
    "    cell_entities = np.arange(num_cells, dtype=np.int32)\n",
    "    pyvista_cells, cell_types = create_vtk_topology(mesh, mesh.topology.dim, cell_entities)\n",
    "    grid = pyvista.UnstructuredGrid(pyvista_cells, cell_types, mesh.geometry.x)\n",
    "    return grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_mesh_plot(mesh):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_mesh_plot(mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Function spaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_velocity = VectorFunctionSpace(mesh, (\"Lagrange\", 2))\n",
    "Y_pressure = FunctionSpace(mesh, (\"Lagrange\", 1))\n",
    "U = VectorFunctionSpace(mesh, (\"Lagrange\", 2))\n",
    "Q_velocity = Y_velocity.clone()\n",
    "Q_pressure = Y_pressure.clone()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Restrictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dofs_Y_velocity = np.arange(0, Y_velocity.dofmap.index_map.size_local + Y_velocity.dofmap.index_map.num_ghosts)\n",
    "dofs_Y_pressure = np.arange(0, Y_pressure.dofmap.index_map.size_local + Y_pressure.dofmap.index_map.num_ghosts)\n",
    "dofs_U = locate_dofs_topological(U, boundaries.dim, boundaries_3)\n",
    "dofs_Q_velocity = dofs_Y_velocity\n",
    "dofs_Q_pressure = dofs_Y_pressure\n",
    "restriction_Y_velocity = DofMapRestriction(Y_velocity.dofmap, dofs_Y_velocity)\n",
    "restriction_Y_pressure = DofMapRestriction(Y_pressure.dofmap, dofs_Y_pressure)\n",
    "restriction_U = DofMapRestriction(U.dofmap, dofs_U)\n",
    "restriction_Q_velocity = DofMapRestriction(Q_velocity.dofmap, dofs_Q_velocity)\n",
    "restriction_Q_pressure = DofMapRestriction(Q_pressure.dofmap, dofs_Q_pressure)\n",
    "restriction = [restriction_Y_velocity, restriction_Y_pressure, restriction_U,\n",
    "               restriction_Q_velocity, restriction_Q_pressure]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Trial and test functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(v, p) = (TrialFunction(Y_velocity), TrialFunction(Y_pressure))\n",
    "(w, q) = (TestFunction(Y_velocity), TestFunction(Y_pressure))\n",
    "u = TrialFunction(U)\n",
    "r = TestFunction(U)\n",
    "(z, b) = (TrialFunction(Q_velocity), TrialFunction(Q_pressure))\n",
    "(s, d) = (TestFunction(Q_velocity), TestFunction(Q_pressure))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### Problem data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nu = 0.04\n",
    "alpha_1 = 0.001\n",
    "alpha_2 = 0.1 * alpha_1\n",
    "x = SpatialCoordinate(mesh)\n",
    "a = 1.0\n",
    "c = 0.8\n",
    "v_d = as_vector((a * (c * 10.0 * (x[1]**3 - x[1]**2 - x[1] + 1.0))\n",
    "                 + ((1.0 - c) * 10.0 * (-x[1]**3 - x[1]**2 + x[1] + 1.0)), 0.0))\n",
    "ff = Constant(mesh, (0., 0.))\n",
    "\n",
    "\n",
    "def g_eval(x):\n",
    "    values = np.zeros((2, x.shape[1]))\n",
    "    values[0, :] = 10.0 * a * (x[1, :] + 1.0) * (1.0 - x[1, :])\n",
    "    return values\n",
    "\n",
    "\n",
    "g = Function(Y_velocity)\n",
    "g.interpolate(g_eval)\n",
    "bc0 = Function(Y_velocity)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimality conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tracking(v, w):\n",
    "    return inner(v, w)(\"-\")\n",
    "\n",
    "\n",
    "def penalty(u, r):\n",
    "    return alpha_1 * inner(grad(u) * t, grad(r) * t) + alpha_2 * inner(u, r)\n",
    "\n",
    "\n",
    "a = [[tracking(v, w) * dS(4), None, None, nu * inner(grad(z), grad(w)) * dx, - b * div(w) * dx],\n",
    "     [None, None, None, - q * div(z) * dx, None],\n",
    "     [None, None, penalty(u, r) * ds(3), - inner(z, r) * ds(3), None],\n",
    "     [nu * inner(grad(v), grad(s)) * dx, - p * div(s) * dx, - inner(u, s) * ds(3), None, None],\n",
    "     [- d * div(v) * dx, None, None, None, None]]\n",
    "f = [tracking(v_d, w) * dS(4),\n",
    "     None,\n",
    "     None,\n",
    "     inner(ff, s) * dx,\n",
    "     None]\n",
    "a[0][0] += Constant(mesh, 0.) * inner(v, w) * (ds(1) + ds(2))\n",
    "a[3][3] = Constant(mesh, 0.) * inner(z, s) * (ds(1) + ds(2))\n",
    "f[1] = Constant(mesh, 0.) * q * dx\n",
    "f[2] = inner(Constant(mesh, (0., 0.)), r) * dx\n",
    "f[4] = Constant(mesh, 0.) * d * dx\n",
    "bdofs_Y_velocity_1 = locate_dofs_topological((Y_velocity, Y_velocity), mesh.topology.dim - 1, boundaries_1)\n",
    "bdofs_Y_velocity_2 = locate_dofs_topological((Y_velocity, Y_velocity), mesh.topology.dim - 1, boundaries_2)\n",
    "bdofs_Q_velocity_12 = locate_dofs_topological((Q_velocity, Y_velocity), mesh.topology.dim - 1, boundaries_12)\n",
    "bc = [DirichletBC(g, bdofs_Y_velocity_1, Y_velocity), DirichletBC(bc0, bdofs_Y_velocity_2, Y_velocity),\n",
    "      DirichletBC(bc0, bdofs_Q_velocity_12, Q_velocity)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(v, p) = (Function(Y_velocity), Function(Y_pressure))\n",
    "u = Function(U)\n",
    "(z, b) = (Function(Q_velocity), Function(Q_pressure))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cost functional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "J = 0.5 * tracking(v - v_d, v - v_d) * dS(4) + 0.5 * penalty(u, u) * ds(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Uncontrolled functional value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract state forms from the optimality conditions\n",
    "a_state = [[replace(a[i][j], {s: w, d: q}) if a[i][j] is not None else None\n",
    "            for j in (0, 1)] for i in (3, 4)]\n",
    "f_state = [replace(f[i], {s: w, d: q}) for i in (3, 4)]\n",
    "bc_state = [bc[0], bc[1]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the block linear system for the state\n",
    "A_state = assemble_matrix_block(a_state, bcs=bc_state,\n",
    "                                restriction=([restriction[i] for i in (3, 4)],\n",
    "                                             [restriction[j] for j in (0, 1)]))\n",
    "A_state.assemble()\n",
    "F_state = assemble_vector_block(f_state, a_state, bcs=bc_state,\n",
    "                                restriction=[restriction[i] for i in (3, 4)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "vp = create_vector_block([f[j] for j in (0, 1)], restriction=[restriction[j] for j in (0, 1)])\n",
    "ksp = PETSc.KSP()\n",
    "ksp.create(mesh.mpi_comm())\n",
    "ksp.setOperators(A_state)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F_state, vp)\n",
    "vp.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the block solution in components\n",
    "with BlockVecSubVectorWrapper(vp, [c.function_space.dofmap for c in (v, p)]) as vp_wrapper:\n",
    "    for vp_wrapper_local, component in zip(vp_wrapper, (v, p)):\n",
    "        with component.vector.localForm() as component_local:\n",
    "            component_local[:] = vp_wrapper_local"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "J_uncontrolled = mesh.mpi_comm().allreduce(assemble_scalar(J), op=MPI.SUM)\n",
    "print(\"Uncontrolled J =\", J_uncontrolled)\n",
    "assert np.isclose(J_uncontrolled, 2.8479865)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_scalar_field_plot(mesh, scalar_field, name):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    grid.point_arrays[name] = scalar_field.compute_point_values()\n",
    "    grid.set_active_scalars(name)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_vector_field_plot(mesh, vector_field, name, factor):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    values = np.zeros((mesh.geometry.x.shape[0], 3))\n",
    "    values[:, :2] = vector_field.compute_point_values()\n",
    "    grid.point_arrays[name] = values\n",
    "    grid.set_active_vectors(name)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    glyphs = grid.glyph(orient=name, factor=factor)\n",
    "    plotter.add_mesh(glyphs)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_vector_field_plot(mesh, v, \"uncontrolled state velocity\", factor=1e-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, p, \"uncontrolled state pressure\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimal control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the block linear system for the optimality conditions\n",
    "A = assemble_matrix_block(a, bcs=bc, restriction=(restriction, restriction))\n",
    "A.assemble()\n",
    "F = assemble_vector_block(f, a, bcs=bc, restriction=restriction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "vpuzb = create_vector_block(f, restriction=restriction)\n",
    "ksp = PETSc.KSP()\n",
    "ksp.create(mesh.mpi_comm())\n",
    "ksp.setOperators(A)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F, vpuzb)\n",
    "vpuzb.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the block solution in components\n",
    "with BlockVecSubVectorWrapper(vpuzb, [c.function_space.dofmap for c in (v, p, u, z, b)],\n",
    "                              restriction) as vpuzb_wrapper:\n",
    "    for vpuzb_wrapper_local, component in zip(vpuzb_wrapper, (v, p, u, z, b)):\n",
    "        with component.vector.localForm() as component_local:\n",
    "            component_local[:] = vpuzb_wrapper_local"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "J_controlled = mesh.mpi_comm().allreduce(assemble_scalar(J), op=MPI.SUM)\n",
    "print(\"Optimal J =\", J_controlled)\n",
    "assert np.isclose(J_controlled, 1.7643950)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_vector_field_plot(mesh, v, \"state velocity\", factor=1e-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, p, \"state pressure\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_vector_field_plot(mesh, u, \"control\", factor=1e-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_vector_field_plot(mesh, z, \"adjoint velocity\", factor=1e-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, b, \"adjoint pressure\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
